# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import datetime
import functools
import logging
import json
import os
import socket
import subprocess
import time
import yaml

from absl import app, flags
from tornado import gen
from tornado.ioloop import IOLoop


def check_utc():
  timestamp = time.time()
  time_now = datetime.datetime.fromtimestamp(timestamp)
  time_utc = datetime.datetime.utcfromtimestamp(timestamp)
  return (time_now - time_utc).total_seconds() == 0


def poll_running_processes(proc_list):
  running_proc_list = []
  for proc in proc_list:
    if proc.poll() is None:
      running_proc_list.append(proc)
    else:
      # Clean up dead process.
      proc.wait(0)
  return running_proc_list


class Scheduler(object):
  def __init__(self, end_time=None):
    self._end_time = end_time

  @property
  def end_time(self):
    return self._end_time

  def update(self, now, warming_up_min, period_min):
    assert warming_up_min >= 0
    assert period_min > 0
    warming_up = datetime.timedelta(minutes=warming_up_min)
    period = datetime.timedelta(minutes=period_min)

    if self._end_time is None or now >= self._end_time - warming_up:
      if (self._end_time is None
          or now >= self._end_time - warming_up + datetime.timedelta(minutes=1)):
        self._end_time = now + warming_up + period
      else:
        self._end_time += period
      return True

    return False


def launch_process(entrypoint, cmd, logpath):
  if isinstance(cmd, list):
    cmd = ' '.join(cmd)

  # Prepare log directory.
  logpath = os.path.normpath(os.path.expanduser(logpath))
  logdir = os.path.dirname(logpath)
  if not os.path.isdir(logdir):
    os.makedirs(logdir, exist_ok=True)

  shell_cmd = './pyrunner %s %s &>> %s' % (entrypoint, cmd, logpath)
  logging.debug('Launch command: %s' % shell_cmd)
  return subprocess.Popen(shell_cmd, shell=True)


@gen.coroutine
def run(argv):
  scheduler = Scheduler()
  proc_list = []

  while True:
    try:
      proc_list = poll_running_processes(proc_list)

      # Read config.
      slot_config = load_slot_config_from_flag()
      if slot_config is None:
        logging.info('No strategy is assigned to slot: %s', get_slot_from_flag())
        continue

      logpath = os.path.normpath(os.path.expanduser(slot_config['log']))
      warming_up_min = int(slot_config['warming_up_min'])
      period_min = int(slot_config['period_min'])
      cmd = [slot_config['cmd'], *argv]

      # Check launchable and launch.
      now = datetime.datetime.now()
      launchable = scheduler.update(now, warming_up_min, period_min)
      if launchable:
        launch_process(slot_config['entrypoint'], cmd, logpath)

        start_time = now + datetime.timedelta(minutes=warming_up_min)
        end_time = start_time + datetime.timedelta(minutes=period_min)
        logging.info('Launched "%s", start trading at %s, end at %s. log="%s"',
                     ' '.join(cmd),
                     start_time.strftime('%Y%m%d %H:%M:%S'),
                     end_time.strftime('%Y%m%d %H:%M:%S'),
                     logpath)

    except Exception:
      logging.exception('Exception occured.')

    finally:
      yield gen.sleep(10)


def load_config(filepath, default_config=None):
  default_config = dict(default_config or {})
  config = None

  for decoder, err_type in [(json, json.decoder.JSONDecodeError), (yaml, yaml.YAMLError)]:
    try:
      with open(filepath, 'r') as f:
        config = decoder.load(f)
        break
    except err_type:
      pass

  if config is None:
    raise ValueError('Invalid config: %s' % filepath)

  if 'default' in config:
    default_config.update(**config['default'])
    del config['default']

  refined = {}
  for slot, slot_config_orig in config.items():
    slot_config = dict(default_config)
    slot_config.update(**slot_config_orig)
    refined[slot] = slot_config

  return refined


def load_config_from_flag():
  FLAGS = flags.FLAGS
  default_config = {}
  if FLAGS.default_strategy_log:
    default_config['log'] = FLAGS.default_strategy_log
  if FLAGS.default_launch_period_min:
    default_config['period_min'] = FLAGS.default_launch_period_min

  config_path = os.path.expanduser(FLAGS.config)
  return load_config(config_path, default_config=default_config)


@functools.lru_cache(maxsize=None)
def get_slot_from_flag():
  slot = flags.FLAGS.slot or socket.getfqdn()
  return slot


def load_slot_config_from_flag():
  slot = get_slot_from_flag()
  config = load_config_from_flag()
  slot_config = config.get(slot, None)
  if (slot_config is None or slot_config.get('cmd', None) is None
      or slot_config.get('disabled', False)):
    return None
  else:
    return slot_config


def main(argv):
  assert check_utc(), 'Invalid timezone'
  logging.info('Slot: %s', get_slot_from_flag())

  try:
    ioloop = IOLoop.current()
    ioloop.add_callback(run, argv[1:])
    ioloop.start()
    return 0

  except KeyboardInterrupt:
    print()
    return 1


if __name__ == '__main__':
  flags.DEFINE_string('config', None, 'Runner config file')
  flags.DEFINE_string('slot', None, 'Runner slot name. Default is FQDN of running host')
  flags.DEFINE_string('default_strategy_log',
                      '~/log/log_default',
                      'Default strategy log file path.')
  flags.DEFINE_integer(
      'default_launch_period_min',
      30,
      'Default strategy launch period. This value will be '
      'used if period is not speicifed in config file.')
  app.run(main)
